// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <turbo/flags/flag.h>
#include <ktest/ktest.h>
#include <kthread/internal/kthread.h>
#include <kthread/internal/condition_variable.h>
#include "kthread/internal/countdown_event.h"
#include <kthread/internal/mutex.h>
#include <kthread/internal/config.h>


int main(int argc, char* argv[]) {
    turbo::set_flag(&FLAGS_task_group_ntags, 3);
    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}

namespace {

std::vector<kthread_tag_t> kutex_wake_return(2, 0);

void* kutex_wake_func(void* arg) {
    auto mutex = static_cast<kthread::Mutex*>(arg);
    kutex_wake_return.push_back(kthread_self_tag());
    mutex->lock();
    kutex_wake_return.push_back(kthread_self_tag());
    mutex->unlock();
    return nullptr;
}

TEST(KthreadButexMultiTest, kutex_wake) {
    kthread::Mutex mutex;
    mutex.lock();
    kthread_t tid1;
    kthread_attr_t attr = KTHREAD_ATTR_NORMAL;
    attr.tag = 1;
    kthread_start_urgent(&tid1, &attr, kutex_wake_func, &mutex);
    mutex.unlock();
    kthread_join(tid1, nullptr);
    ASSERT_EQ(kutex_wake_return[0], kutex_wake_return[1]);
}

std::vector<kthread_tag_t> kutex_wake_all_return1(2, 0);
std::vector<kthread_tag_t> kutex_wake_all_return2(2, 0);

struct ButexWakeAllArgs {
    kthread::CountdownEvent* ev;
    kthread::CountdownEvent* ack;
};

void* kutex_wake_all_func1(void* arg) {
    auto p = static_cast<ButexWakeAllArgs*>(arg);
    auto ev = p->ev;
    auto ack = p->ack;
    kutex_wake_all_return1.push_back(kthread_self_tag());
    ack->signal();
    ev->wait();
    kutex_wake_all_return1.push_back(kthread_self_tag());
    return nullptr;
}

void* kutex_wake_all_func2(void* arg) {
    auto p = static_cast<ButexWakeAllArgs*>(arg);
    auto ev = p->ev;
    auto ack = p->ack;
    kutex_wake_all_return2.push_back(kthread_self_tag());
    ack->signal();
    ev->wait();
    kutex_wake_all_return2.push_back(kthread_self_tag());
    return nullptr;
}

TEST(KthreadButexMultiTest, kutex_wake_all) {
    kthread::CountdownEvent ev(2);
    kthread::CountdownEvent ack(2);
    ButexWakeAllArgs args{&ev, &ack};
    kthread_t tid1, tid2;
    kthread_attr_t attr1 = KTHREAD_ATTR_NORMAL;
    attr1.tag = 1;
    kthread_start_background(&tid1, &attr1, kutex_wake_all_func1, &args);
    kthread_attr_t attr2 = KTHREAD_ATTR_NORMAL;
    attr2.tag = 2;
    kthread_start_background(&tid2, &attr2, kutex_wake_all_func2, &args);
    ack.wait();
    ev.signal(2);
    kthread_join(tid1, nullptr);
    kthread_join(tid2, nullptr);
    ASSERT_EQ(kutex_wake_all_return1[0], kutex_wake_all_return1[1]);
    ASSERT_EQ(kutex_wake_all_return2[0], kutex_wake_all_return2[1]);
}

std::vector<kthread_tag_t> kutex_requeue_return1(2, 0);
std::vector<kthread_tag_t> kutex_requeue_return2(2, 0);

struct ButexRequeueArgs {
    kthread::Mutex* mutex;
    kthread::ConditionVariable* cond;
    kthread::CountdownEvent* ack;
};

void* kutex_requeue_func1(void* arg) {
    auto p = static_cast<ButexRequeueArgs*>(arg);
    auto mutex = p->mutex;
    auto cond = p->cond;
    auto ack = p->ack;
    kutex_wake_all_return1.push_back(kthread_self_tag());
    std::unique_lock<kthread::Mutex> lk(*mutex);
    ack->signal();
    cond->wait(lk);
    kutex_wake_all_return1.push_back(kthread_self_tag());
    return nullptr;
}

void* kutex_requeue_func2(void* arg) {
    auto p = static_cast<ButexRequeueArgs*>(arg);
    auto mutex = p->mutex;
    auto cond = p->cond;
    auto ack = p->ack;
    kutex_wake_all_return2.push_back(kthread_self_tag());
    std::unique_lock<kthread::Mutex> lk(*mutex);
    ack->signal();
    cond->wait(lk);
    kutex_wake_all_return2.push_back(kthread_self_tag());
    return nullptr;
}

TEST(KthreadButexMultiTest, kutex_requeue) {
    kthread::Mutex mutex;
    kthread::ConditionVariable cond;
    kthread::CountdownEvent ack(2);
    ButexRequeueArgs args{&mutex, &cond, &ack};

    kthread_t tid1, tid2;
    kthread_attr_t attr1 = KTHREAD_ATTR_NORMAL;
    attr1.tag = 1;
    kthread_start_background(&tid1, &attr1, kutex_requeue_func1, &args);
    kthread_attr_t attr2 = KTHREAD_ATTR_NORMAL;
    attr2.tag = 2;
    kthread_start_background(&tid2, &attr2, kutex_requeue_func2, &args);
    ack.wait();
    {
        std::unique_lock<kthread::Mutex> lk(mutex);
        cond.notify_all();
    }
    {
        std::unique_lock<kthread::Mutex> lk(mutex);
        cond.notify_all();
    }
    kthread_join(tid1, nullptr);
    kthread_join(tid2, nullptr);
    ASSERT_EQ(kutex_wake_all_return1[0], kutex_wake_all_return1[1]);
    ASSERT_EQ(kutex_wake_all_return2[0], kutex_wake_all_return2[1]);
}

}  // namespace
